"""
fastchat stanford alpaca data convert tools.
"""
import argparse
import json
import pymp
import pathlib

N=56031

def main(data_path, output_path):
    new_data = []
    with pymp.Parallel(192) as p:
        for fno in p.range(1,N+1):
            file_path=pathlib.Path(data_path+'/'+str(fno)+'.json')
            try:
                with file_path.open(encoding = "utf-8") as f:
                    data = json.load(f)
            except Exception as e:
                continue
            i = 1
            n = len(data)
            item = {
                        "type": "chatml",
                        "messages": [
                            {
                                "role": "system",
                                "content": "你是一个专业的心理咨询师。",
                            }
                        ],
                    }
            for example in data:
                if example.get("role", "") == "client":
                    item["messages"].append(
                            {
                                "role": "user",
                                "content": example['content'],
                            }
                    )
                elif example.get("role", "") == "counselor":
                    item["messages"].append(
                            {
                                "role": "assistant",
                                "content": example['content'],
                            }
                    )
                else:
                    print("Error data")
                    return
            new_data.append(item)
    with open(output_path, "w",encoding = "utf-8") as f:
        for sample in new_data:
            f.write(json.dumps(sample,ensure_ascii=False) + '\n')


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_path", type=str, default="path/to/smile/data")
    parser.add_argument(
        "--output_path", type=str, default="output/path/smile-data-conversation.json"
    )
    args = parser.parse_args()
    main(args.data_path, args.output_path)
